import cv2
from PIL import Image
import os
import numpy as np
import torch
from realesrgan import RealESRGANer


def resize_with_opencv(image_path, output_path, scale=2):
    """
    使用OpenCV双三次插值放大图片
    :param image_path: 输入图片路径
    :param output_path: 输出图片路径
    :param scale: 放大倍数（如2表示放大2倍）
    """
    # 读取图片（支持JPG/PNG等格式）
    img = cv2.imread(image_path, cv2.IMREAD_COLOR)  # 读取为BGR格式

    # 获取原始尺寸
    h, w = img.shape[:2]

    # 计算新尺寸（放大scale倍）
    new_w = int(w * scale)
    new_h = int(h * scale)

    # 用双三次插值放大图片（关键参数：cv2.INTER_CUBIC）
    resized_img = cv2.resize(
        img,
        (new_w, new_h),
        interpolation=cv2.INTER_CUBIC,  # 双三次插值（适合放大）
    )

    # 保存结果（可选：转换为RGB格式用于显示）
    cv2.imwrite(output_path, resized_img)
    print(f"放大后的图片已保存至：{output_path}（尺寸：{new_w}x{new_h}）")


# 测试：放大2倍
# resize_with_opencv("小字体票据.png", "resized_by_cv.png", scale=2)


def preprocess_image_for_ocr(
    image_path: str, use_super_resolution: bool = False
) -> Image.Image:
    """
    对输入图像进行一系列预处理，以优化OCR识别效果。
    """
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法使用OpenCV读取图片: {image_path}")

    if use_super_resolution:
        print("正在应用 Real-ESRGAN 超分辨率...")
        img = _run_esrgan(img)

    # 转换为灰度图并进行二值化
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 使用Otsu's方法自动寻找最佳阈值
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # 将处理后的图像从OpenCV格式转换为Pillow格式
    processed_image = Image.fromarray(cv2.cvtColor(binary, cv2.COLOR_GRAY2RGB))
    return processed_image


def _run_esrgan(img: np.ndarray) -> np.ndarray:
    """使用 realesrgan 对图像进行超分辨率处理。"""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"RealESRGAN: Using device: {device}")

    # RealESRGANer会自动处理模型下载和加载
    # model_name='RealESRGAN_x4plus' 是一个通用的高质量模型
    upsampler = RealESRGANer(
        scale=4,
        model_path=os.path.join(
            os.path.dirname(__file__), "4x-ESRGAN.pth"
        ),  # 假设模型在同目录
        model=None,  # 当提供model_path时，此项应为None
        dni_weight=None,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=False,  # 如果GPU支持，可以设为True以加快速度
        gpu_id=None if device.type == "cpu" else torch.cuda.current_device(),
    )

    # 进行超分处理
    output, _ = upsampler.enhance(img, outscale=4)
    return output
